import os
from fastapi import FastAPI, File, UploadFile, Form, Query
from fastapi.responses import JSONResponse, PlainTextResponse
from fastapi.exceptions import HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pathlib import Path
from typing import List
import uvicorn
from fastapi_socketio import SocketManager
from werkzeug.utils import secure_filename
from utils.logger import Logger
logger = Logger()  # 创建logger实例
import time

app = FastAPI()

# 允许所有域名访问所有路由
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 全局变量跟踪模型是否已初始化
selftalk_handler = None
is_model_initialized = False

def get_selftalk_handler():
    """获取 SelfTalkHandler 单例实例，懒加载模式"""
    global selftalk_handler, is_model_initialized
    
    if selftalk_handler is None:
        from SelfTalk import SelfTalkHandler
        logger.info("初始化 SelfTalkHandler 单例（第一次加载可能需要较长时间）")
        selftalk_handler = SelfTalkHandler()
        is_model_initialized = True
    
    return selftalk_handler, is_model_initialized

# 移除服务启动时预加载模型的事件
@app.on_event("startup")
async def startup_event():
    logger.info("服务启动，模型将在首次请求时加载")

# 设置允许上传的文件扩展名
ALLOWED_EXTENSIONS: List[str] = ['mp3', 'wav', 'ogg', 'flac']

def allowed_file(filename: str) -> bool:
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

socket_manager = SocketManager(app=app, mount_location="/ws")

@socket_manager.on('connect',namespace="/ws")
async def connect(sid, environ):
    print('connect connect connect')
    print(sid)
    # print(environ)

@socket_manager.on('message',namespace="/ws")
async def message(sid: str, data: str):
    print(f"Received message from {sid}: {data}")
    await socket_manager.emit('message', data, namespace="/ws")

@app.get("/")
async def home():
    # print("home")
    return {"message": "Flame-render"}

@app.post("/upload-audio")
async def upload_audio(file: UploadFile = File(...)):
    logger.info(f"收到音频上传请求: {file.filename}")
    
    # 检查文件是否为空
    if not file.filename:
        logger.warning("上传失败: 没有选择文件")
        return JSONResponse(status_code=400, content={"message": "No selected file"})

    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        # 保存文件到服务器的磁盘上
        base_path=Path(app.root_path) / 'wav'
        os.makedirs(base_path, exist_ok=True)
        file_location = base_path / filename
        logger.info(f"保存文件到: {file_location}")
        
        try:
            with open(file_location, "wb+") as file_object:
                file_object.write(await file.read())
            logger.info(f"文件 {filename} 保存成功")
            
            # 检查文件是否为wav格式
            file_ext = filename.rsplit('.', 1)[1].lower() if '.' in filename else ''
            if file_ext == 'wav':
                logger.info(f"文件 {filename} 已经是wav格式，无需转换")
                # 对wav文件直接处理，不需要转换格式
                return JSONResponse(status_code=200, content={"message": f'File {filename} uploaded successfully'})
            else:
                # 对非wav文件进行格式转换
                try:
                    wav_npy(file_location) # 转化音频
                    logger.info(f"文件 {filename} 转换成功")
                except Exception as e:
                    logger.error(f"音频转换失败: {str(e)}")
                    return JSONResponse(status_code=500, content={"message": f"Audio conversion failed: {str(e)}"})
                    
                return JSONResponse(status_code=200, content={"message": f'File {filename} uploaded and converted successfully'})
        except Exception as e:
            logger.error(f"文件保存失败: {str(e)}")
            return JSONResponse(status_code=500, content={"message": f"File saving failed: {str(e)}"})
    else:
        logger.warning(f"上传失败: 不支持的文件类型 {file.filename}")
        return JSONResponse(status_code=401, content={"message": 'File type not allowed'})

@app.post("/wav2npy")
async def wav_to_npy(
    file: UploadFile = File(...),
    dataset: str = Form("vocaset")  # 默认数据集
):
    logger.info(f"收到wav2npy请求: {file.filename}, 数据集: {dataset}")
    
    # 检查文件是否为空
    if not file.filename:
        logger.warning("上传失败: 没有选择文件")
        return JSONResponse(status_code=400, content={"message": "No selected file"})

    # 检查文件类型
    if not file.filename.lower().endswith('.wav'):
        logger.warning(f"上传失败: 需要WAV文件，收到: {file.filename}")
        return JSONResponse(status_code=400, content={"message": "Only WAV files are supported"})

    try:
        # 保存文件到服务器的磁盘上
        filename = secure_filename(file.filename)
        base_path = Path(app.root_path) / 'wav'
        os.makedirs(base_path, exist_ok=True)
        file_location = base_path / filename
        logger.info(f"保存文件到: {file_location}")
        
        with open(file_location, "wb+") as file_object:
            file_object.write(await file.read())
        logger.info(f"文件 {filename} 保存成功")
        
        # 获取 SelfTalkHandler 实例并处理文件
        handler, is_first_time = get_selftalk_handler()
        
        # 告知用户首次模型加载可能需要时间
        if is_first_time:
            logger.info("首次加载SelfTalk模型，可能需要较长时间...")
            # 这里可以返回一个初步响应，让用户知道服务器正在加载模型
            # 但由于HTTP请求需要等待完成，这里不中断流程
        
        try:
            start_time = time.time()
            result = handler.process_file(file_location, dataset=dataset)
            process_time = time.time() - start_time
            logger.info(f"文件处理成功: {result['output_path']}, 耗时: {process_time:.2f}秒")
            
            # 将numpy数据转换为JSON格式
            mesh_data = result['mesh_data'].tolist()
            
            # 返回包含mesh_data字段的响应
            response_data = {
                "message": "File processed successfully",
                "output_path": str(result["output_path"]),
                "process_time": f"{process_time:.2f}秒",
                "status": "success",
                "mesh_data": mesh_data
            }
            
            # 如果是第一次加载模型，添加特殊消息
            if is_first_time:
                response_data["first_time_load"] = True
                response_data["message"] = "模型首次加载完成，后续处理将更快"
                global is_model_initialized
                is_model_initialized = False  # 重置为下一次判断
            
            return JSONResponse(status_code=200, content=response_data)
            
        except Exception as e:
            logger.error(f"处理文件失败: {str(e)}")
            return JSONResponse(
                status_code=500, 
                content={
                    "message": f"Failed to process file: {str(e)}",
                    "status": "error"
                }
            )
    except Exception as e:
        logger.error(f"保存文件失败: {str(e)}")
        return JSONResponse(
            status_code=500, 
            content={
                "message": f"Failed to save file: {str(e)}",
                "status": "error"
            }
        )

@app.post("/upload-npy")
async def upload_npy(file: UploadFile = File(...)):
    logger.info(f"收到NPY文件上传请求: {file.filename}")
    
    # 检查文件是否为空
    if not file.filename:
        logger.warning("上传失败: 没有选择文件")
        return JSONResponse(status_code=400, content={"message": "No selected file", "status": "error"})

    # 检查文件类型
    if not file.filename.lower().endswith('.npy'):
        logger.warning(f"上传失败: 需要NPY文件，收到: {file.filename}")
        return JSONResponse(status_code=400, content={"message": "Only NPY files are supported", "status": "error"})

    try:
        # 保存文件到服务器的磁盘上
        filename = secure_filename(file.filename)
        base_path = Path(app.root_path) / 'npy'
        os.makedirs(base_path, exist_ok=True)
        file_location = base_path / filename
        logger.info(f"保存文件到: {file_location}")
        
        with open(file_location, "wb+") as file_object:
            file_object.write(await file.read())
        logger.info(f"文件 {filename} 保存成功")
        
        # 读取NPY文件内容
        try:
            import numpy as np
            mesh_data = np.load(file_location)
            logger.info(f"NPY文件加载成功，形状: {mesh_data.shape}")
            
            # 返回包含mesh_data的响应
            return JSONResponse(
                status_code=200, 
                content={
                    "message": "NPY file uploaded and loaded successfully",
                    "output_path": str(file_location),
                    "status": "success",
                    "mesh_data": mesh_data.tolist()
                }
            )
        except Exception as e:
            logger.error(f"读取NPY文件失败: {str(e)}")
            return JSONResponse(
                status_code=500, 
                content={
                    "message": f"Failed to process NPY file: {str(e)}",
                    "status": "error"
                }
            )
    except Exception as e:
        logger.error(f"保存文件失败: {str(e)}")
        return JSONResponse(
            status_code=500, 
            content={
                "message": f"Failed to save file: {str(e)}",
                "status": "error"
            }
        )

@app.get("/logs/list")
async def get_logs_list():
    """获取所有可用的日志文件列表"""
    try:
        log_dir_path = os.path.join(os.path.dirname(__file__), 'logs')
        if not os.path.exists(log_dir_path):
            return JSONResponse(status_code=200, content={"logs": []})
        
        # 获取所有日志文件
        log_files = []
        for file in os.listdir(log_dir_path):
            if file.startswith("app_") and file.endswith(".log"):
                # 从文件名提取日期 (app_YYYY-MM-DD.log)
                date_str = file[4:-4]  # 去掉 "app_" 和 ".log"
                log_files.append(date_str)
        
        # 按日期排序，最新的日期在前
        log_files.sort(reverse=True)
        
        return JSONResponse(status_code=200, content={"logs": log_files})
    except Exception as e:
        logger.error(f"获取日志列表失败: {str(e)}")
        return JSONResponse(status_code=500, content={"error": f"获取日志列表失败: {str(e)}"})

@app.get("/logs", response_class=PlainTextResponse)
async def get_logs(date: str = None):
    """获取服务器日志内容，可选择特定日期的日志"""
    try:
        log_dir_path = os.path.join(os.path.dirname(__file__), 'logs')
        
        # 如果未指定日期，使用最新的日志文件
        if not date:
            log_file_path = os.path.join(log_dir_path, 'server.log')  # 默认的当前日志
        else:
            log_file_path = os.path.join(log_dir_path, f'app_{date}.log')
        
        if not os.path.exists(log_file_path):
            return f"日志文件不存在: {log_file_path}"
        
        # 读取日志文件内容
        with open(log_file_path, 'r', encoding='utf-8') as f:
            logs = f.read()
        
        return logs
    except Exception as e:
        logger.error(f"获取日志失败: {str(e)}")
        return f"获取日志失败: {str(e)}"

@app.exception_handler(404)
async def not_found_exception(request, exc):
    return JSONResponse(status_code=404, content={"message": "您请求的页面不存在了，请确认后再次访问！"})

if __name__ == '__main__':
    uvicorn.run("server:app", host='127.0.0.1', port=5000, reload=True)